
import numpy as np
import cv2
class EightWayConv():
    """
    根据目标检测的box结果裁切成小图
    args:
        keys: 指定了从data中读取数据的关键字和保存数据的关键字
              格式: keys={
                "in":"your_in_keyword",
                "out":"your_save_keyword"
                        }
    """
    def __init__(self,keys,threshold) -> None:
        self.keys = keys
        self.threshold = threshold


    def __call__(self, data):
        """将需要进行子任务的object裁出
        return:
            data: {img_data, ...}
        """
        seg_res = data[self.keys["in"]]
        res = self.eight_way_conv(seg_res["masks"],self.threshold)
        if self.keys.get("out",False):
            data[self.keys["out"]]["masks"] = res
        else:
            data[self.keys["in"]]["masks"] = res
        return data
    
    @staticmethod
    def eight_way_conv(label_masks, threshold):
        """
        八向卷积, 根据八个方向的像素来判断当前连通域被那些类包围
        """
        for batch_id, label_mask in enumerate(label_masks):
            res_label_mask = label_mask.copy()
            max_y,max_x = label_mask.shape
            max_y-=1
            max_x-=1
            for cls_idx in range(0,label_mask.max()+1):# 不同类分开做连通性检测
                single_cls_result = np.zeros_like(label_mask,dtype=np.uint8)
                single_cls_result[label_mask==cls_idx] = 1
                
                total_area = single_cls_result.sum() # 当前类的总像素数量
                area_threshold = total_area*threshold # 小区域的阈值
                

                # ---1---连通域检测---
                _, labels, stats, centres_xy = cv2.connectedComponentsWithStats(single_cls_result)
                # labels: 标记像素属于哪个连通域, 注意labels==0的表示背景类, stats[x,y,w,h,area]
                # ---2---将小区域的连通域填充为包围其的连通域---
                area_sort = np.argsort(stats[:, -1]) # 从小到大按area排序
                for idx_ in area_sort:
                    if idx_ == 0: # 0为背景类, 不参与计算
                        continue
                    stat = stats[idx_]
                    x1,y1,w,h,area = stat
                    x2 = x1+w
                    y2 = y1+h
                    if (area>(area_threshold)): # 超过阈值则表示至此之后的连通域都超过阈值了
                        break
                    # 八向卷积检测
                    detect_points = [[max(y1-1,0),max(x1-1,0)],     # left_up_point
                                        [(y1+h//2),max(x1-1,0)],         # left_mid_point
                                        [min(y2+1,max_y),max(x1-1,0)], # left_down_point
                                        [max(y1-1,0),(x1+w//2)],       # mid_up_point
                                        [min(y2+1,max_y),(x1+w//2)],   # mid_down_point
                                        [max(y1-1,0),min(x2+1,max_x)], # right_up_point
                                        [(y1+h//2),min(x2+1,max_x)], # right_mid_point
                                        [min(y2+1,max_y),min(x2+1,max_x)] # right_down_point
                                        ]
                    check_values = []
                    for y_,x_ in detect_points:
                        check_values.append(label_mask[y_,x_])
                    max_count_idx = max(set(check_values),key=check_values.count)
                    # 填充类别
                    res_label_mask[labels==idx_] = max_count_idx
            label_masks[batch_id] = res_label_mask
        return label_masks